import argparse
import torch
import torch.utils
import torch.utils.data
from train import train
from utils import init_rand_seed, define_device
init_rand_seed()


BATCH_SIZE = 256
IN_CHANNELS = 3
NUM_CLASSES = 10


if __name__ == '__main__':

    # 获取参数
    parser = argparse.ArgumentParser()
    parser.add_argument('--print_model', action='store_true', default=False, help='print model structure, output shape and params')
    parser.add_argument('--train', action='store_true', default=False)
    parser.add_argument('--eval', action='store_true', default=False)
    parser.add_argument('--model', type=str, default='SimpleNet')
    parser.add_argument('--epochs', type=int, default=3)
    parser.add_argument('--batch-size', type=int, default=128)
    parser.add_argument('--lr', type=float, default=None)
    parser.add_argument('--momentum', type=float, default=0.9)
    parser.add_argument('--weight-decay', type=float, default=2e-5)
    parser.add_argument('--sched', type=str, default=None)
    parser.add_argument('--lr-period', type=int, default=2)
    parser.add_argument('--lr-decay', type=float, default=0.9)
    parser.add_argument('--output', type=str, default=None)
    parser.add_argument('--checkpoint', type=str, default=None)
    args = parser.parse_args()
    print(args)

    device = define_device()
    print('Run on device', device)

    # 启动训练
    train(args, device=device, in_channels=IN_CHANNELS, num_classes=NUM_CLASSES)

